import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from numpy.random import rand
import matplotlib.pyplot as plt
from copy import deepcopy


def Sphere(data):
    total = 0
    data = data ** 2
    for i in data:
        total += i
    return total


def test_func():
    fig = plt.figure()
    ax = Axes3D(fig)
    X = np.arange(-2, 2, 0.05)
    Y = np.arange(-2, 2, 0.05)
    X, Y = np.meshgrid(X, Y)
    Z = np.sin(np.sqrt(X ** 2 + Y ** 2)) / np.sqrt(X ** 2 + Y ** 2) + np.exp(
        (np.cos(2 * np.pi * X) + np.cos(2 * np.pi * Y)) / 2) - 2.71289
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='hot')
    ax.grid(False)


def Sphere1():
    fig = plt.figure()
    ax = Axes3D(fig)
    X = np.arange(-2, 2, 0.05)
    Y = np.arange(-2, 2, 0.05)
    X, Y = np.meshgrid(X, Y)
    Z = (X ** 2 + Y ** 2)
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='hot')
    ax.grid(False)


def Schaffer():
    fig = plt.figure()
    ax = Axes3D(fig)
    X = np.arange(-2, 2, 0.05)
    Y = np.arange(-2, 2, 0.05)
    X, Y = np.meshgrid(X, Y)
    Z = 0.5 - (np.sin(np.sqrt(X ** 2 + Y ** 2)) ** 2 - 0.5) / (1 + 0.001 * (X ** 2 + Y ** 2)) ** 2
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='hot')
    ax.grid(False)


class PSO(object):
    'popsize--------> 种群规模\
    maxgen----------> 最大迭代次数\
    dim-------------> 优化维度\
    popmin,popmax---> 搜寻限值\
    c1,c2-----------> 动态因子\
    vmax vmin-------> 速度限值\
    w1--------------> 权值 '

    def __init__(self, popsize, maxgen, dim, popmin, popmax, c1, c2, w1, vmax, vmin):
        self.popsize = popsize
        self.maxgen = maxgen
        self.dim = dim
        self.popmin = popmin
        self.popmax = popmax
        self.c1 = c1
        self.c2 = c2
        self.w1 = w1
        self.vmax = vmax
        self.vmin = vmin

    def init_population(self):
        pop = rand(self.popsize, self.dim) * (self.popmax - self.popmin)
        v = rand(self.popsize, self.dim) * (self.vmax - self.vmin)
        return pop, v

    def fitness(self, pop):
        fitness_1 = Sphere(pop)
        return fitness_1

    def iter_optimize(self):
        pop, v = PSO.init_population(self)
        current_fit = rand(self.popsize, 1)
        pop_cbest = deepcopy(pop)
        pop_gbest = rand(1, self.dim)
        fit_record = [None] * self.maxgen

        for j in range(self.popsize):
            current_fit[j, :] = PSO.fitness(self, pop_cbest[j, :])

        fit_cbest = deepcopy(current_fit)
        fit_gbest = deepcopy(np.min(fit_cbest))
        bestindex = np.argmin(fit_cbest)
        pop_gbest = pop_cbest[bestindex, :]

        for i in range(self.maxgen):
            for j in range(self.popsize):
                v[j, :] = self.w1 * v[j, :] + self.c1 * rand(1, self.dim) * (
                        pop_cbest[j, :] - pop[j, :]) + self.c2 * rand(1, self.dim) * (pop_gbest - pop[j, :])
                v[j, v[j, :] > self.vmax] = self.vmax
                v[j, v[j, :] < self.vmin] = self.vmin
                pop[j, :] = pop[j, :] + v[j, :]
                pop[j, pop[j, :] > self.popmax] = self.popmax
                pop[j, pop[j, :] < self.popmin] = self.popmin
                current_fit[j, :] = PSO.fitness(self, pop[j, :])

                if current_fit[j, :] < fit_cbest[j, :]:
                    fit_cbest[j, :] = current_fit[j, :]
                    pop_cbest[j, :] = pop[j, :]

                if current_fit[j, :] < fit_gbest:
                    fit_gbest = deepcopy(current_fit[j, :])
                    pop_gbest[:] = pop[j, :]

            fit_record[i] = fit_gbest
        return fit_record

    def drawpicrure(self, data):
        plt.plot(data)
        plt.xlabel('iter_num')
        plt.ylabel('fitness')
        plt.show()

